#!/usr/bin/env python

"""
#2023-07-31: add method: write_band_cover_map
@2023-07-30: finished the initial version.

TODO:
[ ] add support to output coverage area of each band for a specified visit number.
[ ] add member method to generate annimation of coverage of the survey (for different coverage requirements)
[ ] add *ini parsering support
[ ] add *ini configuration output when using non-ini usage
"""

import os, sys
import matplotlib.pylab as plt
import numpy as np
import healpy as hp
import datetime as dt
import time as time
import iniconfig as iniconfig
from tqdm import trange
from cosat.ccd import *

class ObsCoverageStats:
    def __init__(self, nside, debug_mode=False, query_method=0, query_inclusive=False, healpix_map_yrs=None, \
        write_band_cover_map=False, outdir=None):
        # 设置Healpix参数
        self.NSIDE = nside
        self.NPIX = hp.nside2npix( self.NSIDE )

        # 初始话各片探测器的对应的波段和相对焦面中心的位置
        self.bands  = init_ccd_bands()
        self.ccdPos = init_ccd_pos_in_focus()

        # 波段数
        self.num_of_bands = BAND_Y - BAND_NUV + 1

        # 初始化visits统计计数器
        self.wide_visit_counter = np.zeros( (self.num_of_bands, self.NPIX), dtype = np.uint16 )
        self.deep_visit_counter = np.zeros( (self.num_of_bands, self.NPIX), dtype = np.uint16 )

        # 设置各波段覆盖次数要求
        self.wide_visits     = 2
        self.wide_visits_NUV = 4
        self.wide_visits_y   = 4

        self.deep_visits     = 8
        self.deep_visits_NUV = 16
        self.deep_visits_y   = 16

        # 将存储巡天观测指向的数据初始化为None
        self.has_obs_sqe= False
        self.obs_jd     = None
        self.obs_yr     = None
        self.obs_lon    = None
        self.obs_lat    = None

        # 设置搜索方法
        self.query_use_polygon = 0
        self.query_use_disc = 1
        self.query_method = self.query_use_polygon # default
        self.query_method = query_method # default
        self.query_inclusive = query_inclusive

        # 保存部分结果
        self.healpix_map_yrs = []
        self.healpix_map_yrs_flt = []
        if healpix_map_yrs is not None:
            for i in range(len(healpix_map_yrs)):
                self.healpix_map_yrs.append(float(healpix_map_yrs[i]))

        # 输出目录
        self.outdir = './tmp'
        if outdir is not None:
            self.outdir = outdir

        # print('self.healpix_map_yrs {}'.format(self.healpix_map_yrs))
        # sys.exit(0)

        self.write_band_cover_map = write_band_cover_map

        # 记录输出的fits文件
        self.healpix_map_fits_cnt = 0
        self.healpix_map_fits = []

        self.healpix_map_band_cnt = 0

        # debug mode
        self.debug_mode = debug_mode

        self.print_config_info()

    def write_yrTags(self):
        fname = self.outdir + '/yrTags.txt'
        fp = open(fname, 'w', encoding='utf-8')
        fp.write('#idx\tYear\n')
        for i in range(len(self.healpix_map_yrs)):
            fp.write('{:>4d}\t{}\n'.format(i+1, self.healpix_map_yrs_flt[i]))
        fp.close()
        print('> {} is written'.format(fname))


    def read_ini(self, iniFile):
        pass

    def print_config_info(self):
        print('\n********************************************')
        print('> summary of the statistics configuration:')
        print('> healpix:nside = {}'.format(self.NSIDE))
        print('> healpix:npix  = {}'.format(self.NPIX))
        print('> # of bands: {}'.format(self.num_of_bands))

        if self.query_method == self.query_use_disc:
            print('> using healpy.query_disc to search pixels')
        elif self.query_method == self.query_use_polygon:
            print('> using healpy.query_polygon to search pixels')

        if self.query_inclusive is True:
            print('> query inclusive : True')
        else:
            print('> query inclusive : False [default]')

        if self.write_band_cover_map:
            print('> band cover maps will be produced.')

        if self.debug_mode is True:
            print('> runnning in debug mode.')

        print('> outdir: {}'.format(self.outdir))
        print('********************************************')

    def load_obs_seq(self, obs_seq_file=None):
        if obs_seq_file is None:
            print('> obs_seq_file is None, quit.')
            sys.exit(0)

        if os.path.isfile(obs_seq_file):
            seq = np.loadtxt(obs_seq_file)
            self.obs_jd  = seq[:,0]        # original Julian dates
            self.obs_lon = seq[:,2]        # longitude in degrees
            self.obs_lat = seq[:,1]     # latitue in degrees
            self.isDeep  = seq[:,14]

            self.obs_yr     = (seq[:,0] - np.min(seq[:,0])) / 365.25    # 以年为单位的观测时间
            self.has_obs_sqe = True
        else:
            print('> failed to load observation sequence form: {}'.format(obs_seq_file))
            sys.exit(0)

    def getTimeTag(self):
        timeTag = dt.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        return timeTag

    def mkdir(self, dirname=None):
        if dirname is None:
            print('> dirname is None, quit')
            sys.exit(0)
        if os.path.isdir(dirname):
            print('# dir {} already exist'.format(dirname))
        else:
            os.system('mkdir '+dirname)
            print('> created dir: {}'.format(dirname))

    def write_fits(self, idx):
        """
        在输出累计至某运行年份的巡天覆盖面积统计时，将fits文件输出到outdir目录下，
        文件名由运行该程序时的日期+时间（yrTag）来确定
        """
        ofitsname = self.outdir + '/NSIDE_' + str(self.NSIDE) + '_' + str(idx) + '.fits'
        ipix_wide_NUV   = self.wide_visit_counter[0,:] >= self.wide_visits
        # ipix_wide_NUV   = self.wide_visit_counter[0,:] >= self.wide_visits_NUV
        ipix_wide_u     = self.wide_visit_counter[1,:] >= self.wide_visits
        ipix_wide_g     = self.wide_visit_counter[2,:] >= self.wide_visits
        ipix_wide_r     = self.wide_visit_counter[3,:] >= self.wide_visits
        ipix_wide_i     = self.wide_visit_counter[4,:] >= self.wide_visits
        ipix_wide_z     = self.wide_visit_counter[5,:] >= self.wide_visits
        ipix_wide_Y     = self.wide_visit_counter[6,:] >= self.wide_visits
        ipix_wide       = ipix_wide_NUV * ipix_wide_u * ipix_wide_g * ipix_wide_g \
                            * ipix_wide_r * ipix_wide_i * ipix_wide_z * ipix_wide_Y
        # 统计观测次数
        hp.write_map(ofitsname, self.wide_visit_counter[0,:].astype(np.int32), overwrite=True)

        # 统计是否被完整覆盖
        # hp.write_map(ofitsname, ipix_wide.astype(np.int32), overwrite=True)
        # print('* {} has been written.'.format(ofitsname))
    
    def write_band_cover_map_fits(self, idx):
        """
        输出不同波段的覆盖图
        """
        for i in range(len(CCD_bandnames)):
            ofitsname = self.outdir + '/' + CCD_bandnames[i] + '_NSIDE_' + str(self.NSIDE) + '_' + str(idx) + '.fits'
            hp.write_map(ofitsname, self.wide_visit_counter[i,:].astype(np.int32), overwrite=True)
            print('* {} has been written.'.format(ofitsname))

    def get_covered_area(self):
        """
        TODO: 增加其他不同类型的面积统计，如任一波段至少覆盖1次的面积等
        """
        
        ipix_wide_NUV   = self.wide_visit_counter[0,:] >= self.wide_visits_NUV
        ipix_wide_u     = self.wide_visit_counter[1,:] >= self.wide_visits
        ipix_wide_g     = self.wide_visit_counter[2,:] >= self.wide_visits
        ipix_wide_r     = self.wide_visit_counter[3,:] >= self.wide_visits
        ipix_wide_i     = self.wide_visit_counter[4,:] >= self.wide_visits
        ipix_wide_z     = self.wide_visit_counter[5,:] >= self.wide_visits
        ipix_wide_Y     = self.wide_visit_counter[6,:] >= self.wide_visits
        ipix_wide       = np.sum(ipix_wide_NUV * ipix_wide_u * ipix_wide_g * ipix_wide_g \
                            * ipix_wide_r * ipix_wide_i * ipix_wide_z * ipix_wide_Y )
        area_wide = ipix_wide / self.NPIX * 4 * np.pi * (180/np.pi)**2

        ipix_deep_NUV   = self.deep_visit_counter[0,:] >= self.deep_visits_NUV
        ipix_deep_u     = self.deep_visit_counter[1,:] >= self.deep_visits
        ipix_deep_g     = self.deep_visit_counter[2,:] >= self.deep_visits
        ipix_deep_r     = self.deep_visit_counter[3,:] >= self.deep_visits
        ipix_deep_i     = self.deep_visit_counter[4,:] >= self.deep_visits
        ipix_deep_z     = self.deep_visit_counter[5,:] >= self.deep_visits
        ipix_deep_Y     = self.deep_visit_counter[6,:] >= self.deep_visits
        ipix_deep       = np.sum(ipix_deep_NUV * ipix_deep_u * ipix_deep_g * ipix_deep_g \
                            * ipix_deep_r * ipix_deep_i * ipix_deep_z * ipix_deep_Y )
        area_deep = ipix_deep / self.NPIX * 4 * np.pi * (180/np.pi)**2

        return area_wide, area_deep

    def run_stats(self):

        t_start = time.time()

        yrtag = self.getTimeTag()
        self.mkdir(self.outdir)

        if self.debug_mode is True:
            obs_seq_len = 10000 # for debug
        else:
            obs_seq_len = len(self.obs_jd)

        output_stepsize = obs_seq_len // 120 # 控制输出面积统计结果的步长

        fits_cnt = 0
        print('\n> start runnning covered area statistics')

        if self.query_method == self.query_use_disc:
            for i in range(obs_seq_len):    # 对观测指向进行循环
                # 获取当前望远镜的指向
                lon, lat    = self.obs_lon[i], self.obs_lat[i]
                yr = (self.obs_jd[i]-self.obs_jd[0])/365.25

                # 对探测器位置进行循环判断:
                # 首先用query_disc查询一块小天区内的Healpix坐标，这些坐标包围的区域略大于一片CCD对应的天区；
                # 然后对这些坐标进行循环判断，确认它们是否落在CCD四个角所组成的小矩形的范围内；
                # 如果在其内，则对应的计数器+1
                for j in range( len(self.ccdPos) ):
                    y, z = self.ccdPos[j][0], self.ccdPos[j][1]
                    p1, p2, p3, p4 = computeCCDPos( lon=lon, lat=lat, y=y, z=z )
                    # 获取当前这一片探测器的中心指向，然后缩小搜索范围，面积统计的效率会更高
                    p = (p1+p2+p3+p4) / 4
                    ipix = hp.query_disc( nside=self.NSIDE, vec=p, radius=np.radians(0.175), inclusive=self.query_inclusive )
                    for k in range(len(ipix)):
                        vec = hp.pix2vec(self.NSIDE, ipix[k])
                        if isInRect(vec, p1, p2, p3, p4) is True:
                            self.wide_visit_counter[self.bands[j],ipix[k]] += 1
                            self.deep_visit_counter[self.bands[j],ipix[k]] += (self.isDeep[i] > 0)

                if i % output_stepsize == 0:
                    area_wide, area_deep = self.get_covered_area()
                    print('> Obs date = {:.4f} Yr, covered area: wide_field = {:>8.3f} deg^2, deep_field = {:>6.3f} deg^2'.format(yr, area_wide, area_deep))

                if len(self.healpix_map_yrs) > 0:
                    if fits_cnt < len(self.healpix_map_yrs) and yr > self.healpix_map_yrs[fits_cnt]:
                        yrtag = 'YEAR' + '-' + str(self.healpix_map_yrs[fits_cnt])
                        self.write_fits(fits_cnt+1)
                        if self.write_band_cover_map is True:
                            self.write_band_cover_map_fits(fits_cnt+1)
                        fits_cnt += 1
        elif self.query_method == self.query_use_polygon:
            yr = None
            for i in range(obs_seq_len):    # 对观测指向进行循环
                # 获取当前望远镜的指向
                lon, lat    = self.obs_lon[i], self.obs_lat[i]
                yr = (self.obs_jd[i]-self.obs_jd[0])/365.25

                # 对探测器位置进行循环判断：
                # 首先用query_polygon直接查询落在CCD四个角所组成的polygon范围内的Healpix坐标；
                # 然后将对应Healpix的计数器进行更新
                for j in range( len(self.ccdPos) ):
                    y, z = self.ccdPos[j][0], self.ccdPos[j][1]
                    p1, p2, p3, p4 = computeCCDPos( lon=lon, lat=lat, y=y, z=z )
                    polygon = np.vstack((p1,p2,p3,p4))
                    ipix = hp.query_polygon( nside=self.NSIDE, vertices=polygon, inclusive=self.query_inclusive )
                    self.wide_visit_counter[self.bands[j],ipix] += 1
                    self.deep_visit_counter[self.bands[j],ipix] += (self.isDeep[i] > 0)

                if i % output_stepsize == 0:
                    area_wide, area_deep = self.get_covered_area()
                    print('> Obs date = {:.4f} Yr, covered area: wide_field = {:>8.3f} deg^2, deep_field = {:>6.3f} deg^2'.format(yr,area_wide, area_deep))

                if len(self.healpix_map_yrs) > 0:
                    if fits_cnt < len(self.healpix_map_yrs) and yr > self.healpix_map_yrs[fits_cnt]:
                        yrtag = 'YEAR' + '-' + str(self.healpix_map_yrs[fits_cnt])
                        self.write_fits(fits_cnt+1)
                        if self.write_band_cover_map is True:
                            self.write_band_cover_map_fits(fits_cnt+1)
                        self.healpix_map_yrs_flt.append(round(yr,3))
                        fits_cnt += 1
        
            self.write_fits(fits_cnt+1)
            if self.write_band_cover_map is True:
                self.write_band_cover_map_fits(fits_cnt+1)
            self.healpix_map_yrs_flt.append(round(yr,3))

        # 全部观测指向统计之后
        self.write_fits('ALL')

        t_end = time.time()
        seconds = t_end - t_start
        m, s = divmod(seconds, 60)
        h, m = divmod(m, 60)
        print('\n> Total time used: %02d:%02d:%02d'%(h, m, s))
        self.write_yrTags()

        

    def print_res(self):
        pass
